#include "eigen/net.h"

namespace  ldl_eigen
{
void Net::add_hidden_layer(std::shared_ptr<HiddenLayer> layer)
{
    if(hidden_layers_.size() > 0)
    {
        layer->set_input(hidden_layers_.at(hidden_layers_.size() - 1)->output());
        hidden_layers_.at(hidden_layers_.size() - 1)->set_output_gradient(layer->input_gradient());
    }
    hidden_layers_.push_back(layer);
}

void Net::set_loss(std::shared_ptr<Loss> loss)
{
    loss_ = loss;
}

float Net::train(const Eigen::MatrixXf &features, const Eigen::MatrixXf &labels)
{
    hidden_layers_.at(0)->set_input(features);
    
    for(auto &hidden_layer : hidden_layers_)
    {
        hidden_layer->forward();
    }
    auto loss = loss_->forward(hidden_layers_.at(hidden_layers_.size() - 1)->output(), labels);

    auto input_gradient = loss_->backward();
    hidden_layers_.at(hidden_layers_.size() - 1)->set_output_gradient(input_gradient);
    for(int index = hidden_layers_.size() - 1;index >= 0 ;index--)
    {
        hidden_layers_.at(index)->backward();
        hidden_layers_.at(index)->update();
    }
    return loss;
}

Eigen::MatrixXf Net::predict(const Eigen::MatrixXf &features)
{
    hidden_layers_.at(0)->set_input(features);
    
    for(auto &hidden_layer : hidden_layers_)
    {
        hidden_layer->forward();
    }
    return hidden_layers_.at(hidden_layers_.size() - 1)->output();
}
}
